package com.wd.mqtt;

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousServerSocketChannel;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.nio.charset.StandardCharsets;
import java.util.concurrent.Executors;

public class MqttAioVirtualThreadServer {

    private static final int PORT = 1883;

    public static void main(String[] args) throws IOException {
        var executor = Executors.newVirtualThreadPerTaskExecutor();

        try (AsynchronousServerSocketChannel serverSocket = AsynchronousServerSocketChannel.open()) {
            serverSocket.bind(new InetSocketAddress(PORT));
            System.out.println("MQTT server started on port " + PORT);

            serverSocket.accept(null, new CompletionHandler<AsynchronousSocketChannel,Void>() {
                @Override
                public void completed(AsynchronousSocketChannel client, Void attachment) {
                    // Accept next connection
                    try {
                        client.setOption(java.net.StandardSocketOptions.SO_RCVBUF, 1024*8);
                        client.setOption(java.net.StandardSocketOptions.SO_SNDBUF, 1024*8);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }

                    serverSocket.accept(null, this);
                    // Handle this connection in a virtual thread
                    executor.submit(() -> handleClient(client));
                }

                @Override
                public void failed(Throwable exc, Void attachment) {
                    exc.printStackTrace();
                }
            });

            // Prevent main thread from exiting
            Thread.currentThread().join();
        } catch (Exception e) {
            e.printStackTrace();
        }
    }

    private static void handleClient(AsynchronousSocketChannel client) {
        ByteBuffer buffer = ByteBuffer.allocate(20);

//        client.read();
        client.read(buffer, buffer, new CompletionHandler<>() {
            @Override
            public void completed(Integer result, ByteBuffer buffer) {
                if (result == -1) {
                    try {
                        client.close();
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                    return;
                }

                buffer.flip();
                byte messageType = (byte) (buffer.get() >> 4);
                byte remainingLength = buffer.get();

                if (messageType == 1) { // MQTT CONNECT
                    try {
                        handleConnect(buffer, client);
                    } catch (IOException e) {
                        e.printStackTrace();
                    }
                }



                buffer.clear();

                client.read(buffer, buffer, this);
            }

            @Override
            public void failed(Throwable exc, ByteBuffer buffer) {
                try {
                    client.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        });
    }

    private static void handleConnect(ByteBuffer buffer, AsynchronousSocketChannel client) throws IOException {
        // Read the protocol name length
        short protocolNameLength = buffer.getShort();
        byte[] protocolNameBytes = new byte[protocolNameLength];
        buffer.get(protocolNameBytes);
        String protocolName = new String(protocolNameBytes, StandardCharsets.UTF_8);

        // Read protocol level and connect flags
        byte protocolLevel = buffer.get();
        byte connectFlags = buffer.get();

        // Read the keep alive
        short keepAlive = buffer.getShort();

        // Respond with CONNACK
        ByteBuffer connackBuffer = ByteBuffer.allocate(4);
        connackBuffer.put((byte) 0x20); // CONNACK type
        connackBuffer.put((byte) 0x02); // Remaining length
        connackBuffer.put((byte) 0x00); // Connect acknowledge flags
        connackBuffer.put((byte) 0x00); // Connect return code (0: Connection Accepted)
        connackBuffer.flip();

        client.write(connackBuffer, connackBuffer, new CompletionHandler<>() {
            @Override
            public void completed(Integer result, ByteBuffer buffer) {
                // Response sent, do nothing or prepare for next communication
            }

            @Override
            public void failed(Throwable exc, ByteBuffer buffer) {
                try {
                    client.close();
                } catch (IOException e) {
                    e.printStackTrace();
                }
            }
        });
    }
}
